package trainer;

import java.util.*;
import java.io.*;

/**
 * 
 * @author Chung-hau Liang
 *
 */
public class TrainModelByFile {

  HashMap<String, Double> unigram_count = new HashMap<String, Double>();

  HashMap<String, Double> bigram_count = new HashMap<String, Double>();

  HashMap<String, Double> trigram_count = new HashMap<String, Double>();

  HashMap<String, Double> unigram_prob = new HashMap<String, Double>();

  HashMap<String, Double> bigram_prob = new HashMap<String, Double>();

  HashMap<String, Double> trigram_prob = new HashMap<String, Double>();

  HashMap<String, Double> unigram_de = new HashMap<String, Double>();

  HashMap<String, Double> bigram_de = new HashMap<String, Double>();

  ArrayList<String> train_data = new ArrayList<String>();

  int nUnigram = 0, nBigram = 0, nTrigram = 0;

  double lambda = 0.2;

  double dNumTag = 35.0;

  int nX = 35;

  double dUniPP = 0.0, dBiPP = 0.0, dTriPP = 0.0;

  StaticModel sm = null;

  /**
   * @param args
   *          [0]: filename of training data;
   * @param args
   *          [1]: filename of model.
   */
  public static void main(String[] args) {

    // find best lambda for model adaption.
    // HeldOutTest hot = new HeldOutTest();
    // double dL = hot.HeldOutTest(args[0]);//return the best lambda

    // Start to build model given best lambda.
    TrainModelByFile tm = new TrainModelByFile("TrainA.txt", "TrainA_model.txt", 0.2);
    // The output will be a model file named by args[1].
  }

  /**
   * Constructor
   */
  public TrainModelByFile(String filename_train, String filename_model, double dL) {

    // Set lambda
    lambda = dL;
    // Read in the training data
    if (!loadData(filename_train)) {
      System.err.println("Failed to read training data :" + filename_train);
    }

    // Load static model
    sm = new StaticModel();
    sm.loadRules("no_rule.txt");
    // Counting unigram
    countUnigram();

    // Counting bigram
    countBigram();

    // Count trigram
    countTrigram();

    // output
    outputModel(filename_model);
    System.err.println("Done");

  }

  /**
   * Load data from filename
   * 
   * @param filename
   * @return
   */
  private boolean loadData(String filename) {
    boolean bOk = false;
    try {
      BufferedReader in = new BufferedReader(new FileReader(filename));
      String line = "";
      train_data.add("<S>");
      train_data.add("<S>");
      while ((line = in.readLine()) != null) {
        train_data.add(line);
        if (line.equalsIgnoreCase("<PERIOD>")) {
          train_data.add("<S>");
          train_data.add("<S>");
        }
      }
      in.close();
      bOk = true;
    } catch (IOException e) {
    }
    return bOk;
  }

  /**
   * Count the unigram
   */
  private void countUnigram() {

    // Add-Lambda Smoothing
    for (String s : train_data) {
      if ((!s.equals("<S>")))
        nUnigram++;

      if (unigram_count.get(s) == null) {
        unigram_count.put(s, 1.0);// switch to 1 for add one
      } else {
        double nC = unigram_count.get(s).doubleValue();
        nC++;
        unigram_count.put(s, nC);
      }
    }

  }

  /**
   * Count the bigram
   */
  private void countBigram() {
    String prev = null;
    for (String current : train_data) {
      if (prev != null) {
        String key = prev + ' ' + current;
        Double count = bigram_count.get(key);
        if (count == null) {
          count = 0.0;
        }

        count++;

        bigram_count.put(key, count);
      }

      prev = current;
    }
  }

  /**
   * Count the trigram
   */
  private void countTrigram() {
    trigram_count = countNgram(3);
  }

  /**
   * General Function to count N-gram
   * 
   * @param n
   * @return
   */
  private HashMap<String, Double> countNgram(int n) {
    HashMap<String, Double> counts = new HashMap<String, Double>();
    LinkedList<String> history = new LinkedList<String>();

    for (String current : train_data) {
      if (history.size() >= 2) {
        String key = "";
        for (String historyTerm : history) {
          key += historyTerm + ' ';
        }
        key += current;

        Double count = counts.get(key);
        if (count == null) {
          count = 0.0;
        }

        count++;

        counts.put(key, count);
      }

      history.addLast(current);
      if (history.size() > 2) {
        history.removeFirst();
      }
    }

    return counts;
  }

  /**
   * Generate the N-gram model
   * 
   * @param filename
   */
  private void outputModel(String filename) {
    try {
      PrintWriter out = new PrintWriter(new FileWriter(filename));

      // Unigram
      double dSum = 0.0;
      out.println("Unigram:");

      for (String tag : Vocabulary.getTags()) {
        Double nC = unigram_count.get(tag);
        if (nC == null) {
          nC = 0.0;
        }

        double dC = (nC + lambda) / ((double) nUnigram + lambda * Vocabulary.getNumTags());

        // combined score with static_model
        // dC = dC * lambda + (1-lambda) * sm.getStaticModel(tag,1);

        // sum up prob. of each unigram
        dSum += dC;
        unigram_prob.put(tag, dC);
        dUniPP += Math.log(1.0 / dC);
        out.println(tag + "\t" + String.valueOf(dC));
      }
      dUniPP /= (double) unigram_count.size();// cross entropy
      dUniPP = Math.pow(2, dUniPP);// Perplexity
      out.println("Unigram total:" + String.valueOf(dSum) + "\tPerplexity:"
              + String.valueOf(dUniPP));

      // Bigram
      dSum = 0.0;
      out.println("Bigram:");

      String[] historyTags = new String[Vocabulary.getNumTags() + 1];
      int counter = 0;
      for (String tag : Vocabulary.getTags()) {
        historyTags[counter++] = tag;
      }
      historyTags[counter++] = Vocabulary.getStartSentenceTag();

      for (String historyTag : historyTags) {
        // First, count the number of times the history occurred
        // i.e., C(*|historyTag)
        double nU = 0;
        for (String currentTag : Vocabulary.getTags()) {
          String key = historyTag + ' ' + currentTag;
          Double conditionalProb = bigram_count.get(key);
          if (conditionalProb != null) {
            nU += conditionalProb;
          }
        }

        // Next, compute bigram probability
        for (String currentTag : Vocabulary.getTags()) {
          String key = historyTag + ' ' + currentTag;
          Double nC = bigram_count.get(key);
          if (nC == null) {
            nC = 0.0;
          }

          // Prob. of bigram = bigram_count / unigram_count
          double dC = (nC + lambda) / ((double) nU + lambda * Vocabulary.getNumTags());// /(1.0/Math.log(nU+1));//+34.0);
          if (dC == 0)
            dC = 0.00000000000001;
          bigram_prob.put(key, dC);
          // combined score with static_model
          // dC = dC * lambda + (1-lambda) * sm.getStaticModel(tag,2);
          // sum up prob. of each unigram

          dSum += dC;
          dBiPP += Math.log(1.0 / dC);// cross entropy
          out.println(key + "\t" + String.valueOf(dC));
        }
      }
      dBiPP /= (double) bigram_count.size();// Perplexity
      dBiPP = Math.pow(2, dBiPP);
      out.println("Bigram total:" + String.valueOf(dSum) + "\tPerplexity:" + String.valueOf(dBiPP));

      // Trigram
      dSum = 0.0;
      out.println("Trigram:");
      for (String prevPrevTag : historyTags) {
        for (String prevTag : historyTags) {
          String history = prevPrevTag + ' ' + prevTag;

          // First, count the number of times the history occurred
          // i.e., C(*|history)
          double nB = 0;
          for (String currentTag : Vocabulary.getTags()) {
            String key = history + ' ' + currentTag;
            Double conditionalProb = trigram_count.get(key);
            if (conditionalProb != null) {
              nB += conditionalProb;
            }
          }

          // Next, compute the probability
          for (String currentTag : Vocabulary.getTags()) {
            String key = history + ' ' + currentTag;

            Double nC = trigram_count.get(key);
            if (nC == null) {
              nC = 0.0;
            }

            double dC = (nC + lambda) / (nB + lambda * Vocabulary.getNumTags());

            // combined score with static_model
            // dC = dC * lambda + (1-lambda) * sm.getStaticModel(tag,3);
            // sum up prob. of each unigram
            // here is to avoid infinity of perplexity
            if (dC == 0)
              dC = 0.00000000000001;
            trigram_prob.put(key, dC);
            dSum += dC;
            dTriPP += Math.log(1.0 / dC);// cross entropy
            out.println(key + "\t" + String.valueOf(dC));
          }
        }
      }
      dTriPP /= (double) trigram_count.size();// perplexity
      dTriPP = Math.pow(2, dTriPP);
      out.println("Trigram total:" + String.valueOf(dSum) + "\tPerplexity:"
              + String.valueOf(dTriPP));

      out.close();
    } catch (IOException e) {
    }
  }
}